from typing import List, Optional
from src.logging import get_logger
from src.utils.tree_visualizer import TreeVisualizer


class TreeNode:
    """
    多叉树的节点类。
    属性：
        child: 子节点列表
        data: 节点数据，包含token（分块后的token列表）、pre_hash（父节点哈希）、curr_hash（当前节点哈希）
    """
    def __init__(self, token, pre_hash, curr_hash):
        self.child = []  # 存放子节点
        self.data = {
            'token': token,
            'pre_hash': pre_hash,
            'curr_hash': curr_hash
        }

    def __repr__(self):
        return f"TreeNode(token={self.data['token']}, curr_hash={self.data['curr_hash']})"


class IndexTree:
    """
    多叉树结构，支持节点的插入、删除和token匹配深度查询。
    属性：
        root: 根节点
        block_size: token分块大小
    """
    def __init__(self, block_size):
        self.root = None
        self.block_size = block_size
        self.logger = get_logger("index_tree")
        self.visualizer = TreeVisualizer()
        self.logger.info("IndexTree created", block_size=block_size)

    def insert_node(self, token: List[int], pre_hash: Optional[int], curr_hash: int):
        """
        插入节点。
        :param token: 节点的token（list）
        :param pre_hash: 父节点的curr_hash
        :param curr_hash: 当前节点的curr_hash
        """
        if not self.root:
            self.root = TreeNode(token, pre_hash, curr_hash)
            self.logger.debug("Created root node",
                            curr_hash=curr_hash,
                            token_length=len(token) if token else 0)
            
            # DEBUG级别的树状态日志
            debug_log = self.visualizer.generate_tree_debug_log(
                self, "INSERT", curr_hash,
                curr_hash=curr_hash, pre_hash=pre_hash, tokens=token
            )
            if debug_log:
                self.logger.debug(debug_log)
            return
        
        parent = self._find_node_by_curr_hash(self.root, pre_hash)
        if not parent:
            self.logger.error("Parent node not found",
                            pre_hash=pre_hash,
                            curr_hash=curr_hash)
            raise Exception(f"未找到pre_hash为 {pre_hash} 的父节点")
        
        new_node = TreeNode(token, pre_hash, curr_hash)
        parent.child.append(new_node)
        self.logger.debug("Node inserted",
                         pre_hash=pre_hash,
                         curr_hash=curr_hash,
                         token_length=len(token) if token else 0,
                         parent_children_count=len(parent.child))
        
        # DEBUG级别的树状态日志
        debug_log = self.visualizer.generate_tree_debug_log(
            self, "INSERT", curr_hash,
            curr_hash=curr_hash, pre_hash=pre_hash, tokens=token
        )
        if debug_log:
            self.logger.debug(debug_log)

    def delete_node(self, pre_hash: int, curr_hash: int):
        """
        删除节点。
        :param pre_hash: 父节点的curr_hash
        :param curr_hash: 目标节点的curr_hash
        :return: 删除成功返回True，否则抛出异常
        """
        if not self.root:
            self.logger.error("Attempted to delete from uninitialized tree")
            raise Exception("树未初始化")
        
        parent = self._find_node_by_curr_hash(self.root, pre_hash)
        if not parent:
            self.logger.error("Parent node not found for deletion",
                            pre_hash=pre_hash,
                            curr_hash=curr_hash)
            raise Exception(f"未找到pre_hash为 {pre_hash} 的父节点")
        
        for i, node in enumerate(parent.child):
            if node.data['curr_hash'] == curr_hash:
                del parent.child[i]
                self.logger.debug("Node deleted",
                               pre_hash=pre_hash,
                               curr_hash=curr_hash,
                               remaining_children=len(parent.child))
                
                # DEBUG级别的树状态日志
                debug_log = self.visualizer.generate_tree_debug_log(
                    self, "DELETE",
                    curr_hash=curr_hash, pre_hash=pre_hash
                )
                if debug_log:
                    self.logger.debug(debug_log)
                
                return True
        
        self.logger.warn("Node not found for deletion",
                        pre_hash=pre_hash,
                        curr_hash=curr_hash)
        raise Exception("未找到要删除的节点")

    def _find_node_by_curr_hash(self, node: TreeNode, curr_hash: int) -> Optional[TreeNode]:
        """
        递归查找curr_hash对应的节点。
        :param node: 当前遍历到的节点
        :param curr_hash: 需要查找的curr_hash
        :return: 找到则返回节点，否则返回None
        """
        if node.data['curr_hash'] == curr_hash:
            return node
        for child in node.child:
            res = self._find_node_by_curr_hash(child, curr_hash)
            if res:
                return res
        return None

    def find_parent_by_curr_hash(self, curr_hash: int) -> Optional[TreeNode]:
        """
        递归查找curr_hash对应节点的父节点。
        :param curr_hash: 需要查找的子节点curr_hash
        :return: 找到则返回父节点，否则返回None
        """
        return self._find_parent_by_curr_hash(self.root, curr_hash)

    def _find_parent_by_curr_hash(self, node: Optional[TreeNode], curr_hash: int) -> Optional[TreeNode]:
        """
        递归查找curr_hash对应节点的父节点。
        :param node: 当前遍历到的节点
        :param curr_hash: 需要查找的子节点curr_hash
        :return: 找到则返回父节点，否则返回None
        """
        if not node:
            return None
        
        # 检查当前节点的子节点中是否有目标节点
        for child in node.child:
            if child.data['curr_hash'] == curr_hash:
                return node
        
        # 递归查找子树
        for child in node.child:
            res = self._find_parent_by_curr_hash(child, curr_hash)
            if res:
                return res
        return None

    def match_token_size(self, token: List[int]) -> int:
        """
        查询token（列表）在树中匹配的最大深度，返回最大深度*self.block_size。
        :param token: 需要查找的token（列表）
        :return: 最大匹配长度（int）
        """
        if not self.root:
            self.logger.debug("Match attempted on empty tree", token_length=len(token))
            return 0
        
        token_block_list = []
        start_index = 0
        while start_index < len(token):
            token_block: list[int] = token[start_index:start_index + self.block_size]
            token_block_list.append(token_block)
            start_index += self.block_size

        self.logger.debug("Starting token matching",
                        token_length=len(token),
                        block_count=len(token_block_list),
                        block_size=self.block_size)

        def dfs(node, depth):
            if depth >= len(token_block_list):
                return depth
            if node.data['token'] != token_block_list[depth]:
                return depth
            max_depth = depth + 1
            for child in node.child:
                child_depth = dfs(child, depth + 1)
                if child_depth > max_depth:
                    max_depth = child_depth
            return max_depth

        max_depth = dfs(self.root, 0)
        match_size = max_depth * self.block_size
        
        self.logger.debug("Token matching completed",
                        max_depth=max_depth,
                        match_size=match_size,
                        match_ratio=round(match_size / len(token), 2) if len(token) > 0 else 0)
        
        return match_size
